import cv2
import random
import colorsys
import numpy as np
import tensorflow as tf
from core.config import cfg

"""
从文件中读取类名
Arguments:
    class_file_name:文件路径
Return：
    类名,dict类型
"""
def read_class_names(class_file_name):
    names = {}
    with open(class_file_name, 'r') as data:
        for ID, name in enumerate(data):
            names[ID] = name.strip('\n')
    return names

"""
从anchor_path中读取聚类后的anchor box
组成3*3的矩阵，矩阵的每个元素为w h
"""
def get_anchors(anchors_path):
    with open(anchors_path) as f:
        anchors = f.readline()
    anchors = np.array(anchors.split(','), dtype=np.float32)
    return anchors.reshape(3, 3, 2)

"""
图像预处理

"""
def image_preporcess(image, target_size, gt_boxes=None):
    #cv.imread(img_path)读取的是BGR格式
    #将BGR格式转换为RGB格式
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)
    #Darknet53输入尺寸
    ih, iw    = target_size
    #image.shape 返回0，1轴的像素数
    #原始图像大小
    h,  w, _  = image.shape
    """
    依据Darknet53要求的输入图像尺寸（target_size）缩放原始图像的尺寸
    原始图像宽、高中最大值设为输入尺寸，另外一个依比例缩放
    image图像尺寸小于target_size的，会放大一些
    image图像尺寸大于target_size的，会缩小一些
    """
    scale = min(iw/w, ih/h)
    nw, nh  = int(scale * w), int(scale * h)
    image_resized = cv2.resize(image, (nw, nh))
    #缩放后不足原始尺寸的部分用128填充（中灰色）
    image_paded = np.full(shape=[ih, iw, 3], fill_value=128.0)
    dw, dh = (iw - nw) // 2, (ih-nh) // 2
    image_paded[dh:nh+dh, dw:nw+dw, :] = image_resized
    #图像归一化
    image_paded = image_paded / 255.

    if gt_boxes is None:
        return image_paded

    else:
        gt_boxes[:, [0, 2]] = gt_boxes[:, [0, 2]] * scale + dw
        gt_boxes[:, [1, 3]] = gt_boxes[:, [1, 3]] * scale + dh
        return image_paded, gt_boxes

"""
显示bboxes
Argument 
    image:原始图像:从屏幕截取的图像
    bboxes:经过nms处理后的bboxes
    classes:类别名称
    show_lable:是否显示bboxes的类别和confidence信息
Return 
    image:包含bboxes的图像信息
"""
def draw_bbox(image, bboxes, classes=read_class_names(cfg.YOLO.CLASSES), show_label=True):
    """
    bboxes: [x_min, y_min, x_max, y_max, probability, cls_id] format coordinates.
    """
    num_classes = len(classes)
    image_h, image_w, _ = image.shape
    #对每个类别的bbox分配一个颜色为hsv模式即色相、亮度、饱和度
    hsv_tuples = [(1.0 * x / num_classes, 1., 1.) for x in range(num_classes)]
    #将hsv模式转换为rgb模式
    colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
    colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), colors))

    #打乱颜色
    random.seed(0)
    random.shuffle(colors)
    random.seed(None)

    #显示每一个bbox
    for i, bbox in enumerate(bboxes):
        coor = np.array(bbox[:4], dtype=np.int32)
        fontScale = 0.5
        score = bbox[4]
        class_ind = int(bbox[5])
        #bbox的颜色
        bbox_color = colors[class_ind]
        #bbox的线宽
        bbox_thick = int(0.6 * (image_h + image_w) / 600)
        #画出矩形
        #TODO:需要修改这个源代码
        c1, c2 = (coor[0], coor[1]), (coor[2], coor[3])
        cv2.rectangle(image, c1, c2, bbox_color, bbox_thick)

        if show_label:
            bbox_mess = '%s: %.2f' % (classes[class_ind], score)
            #计算字符串的高度和宽度
            t_size = cv2.getTextSize(bbox_mess, 0, fontScale, thickness=bbox_thick//2)[0]
            #标签信息的矩形
            cv2.rectangle(image, c1, (c1[0] + t_size[0], c1[1] - t_size[1] - 3), bbox_color, -1)  # filled
            #显示类别和置信度信息
            cv2.putText(image, bbox_mess, (c1[0], c1[1]-2), cv2.FONT_HERSHEY_SIMPLEX,
                        fontScale, (0, 0, 0), bbox_thick//2, lineType=cv2.LINE_AA)

    return image


"""
计算boxes1和boxes2的IOU值

Argument
    bboxes1,bboxes2:bounding box    
Return 
    ious:bboxes1和bboxes2的iou值
"""

def bboxes_iou(boxes1, boxes2):

    boxes1 = np.array(boxes1)
    boxes2 = np.array(boxes2)

    boxes1_area = (boxes1[..., 2] - boxes1[..., 0]) * (boxes1[..., 3] - boxes1[..., 1])
    boxes2_area = (boxes2[..., 2] - boxes2[..., 0]) * (boxes2[..., 3] - boxes2[..., 1])

    left_up       = np.maximum(boxes1[..., :2], boxes2[..., :2])
    right_down    = np.minimum(boxes1[..., 2:], boxes2[..., 2:])

    inter_section = np.maximum(right_down - left_up, 0.0)
    inter_area    = inter_section[..., 0] * inter_section[..., 1]
    union_area    = boxes1_area + boxes2_area - inter_area
    ious          = np.maximum(1.0 * inter_area / union_area, np.finfo(np.float32).eps)

    return ious

"""
从pb文件中获取tensor
Argument
    graph:默认的计算图
    pb_file:pb文件路径
    return_elements:需要从pb_file计算图模型中加载到默认计算图模型中的tensor或op
Return
    加载后的tensor
"""
def read_pb_return_tensors(graph, pb_file, return_elements):
    #打开计算图文件
    with tf.gfile.FastGFile(pb_file, 'rb') as f:
        """
        tf.GraphDef()用于操作计算图模型
        f.read()读取计算图pb文件中的内容，以str类型返回
        """
        frozen_graph_def = tf.GraphDef()
        frozen_graph_def.ParseFromString(f.read())
    #定义graph为当前线程的默认计算图模型
    with graph.as_default():
        """
        将计算图从frozen_graph_def导入到当前默认计算图中
        frozen_graph_def中包含了需要加载到默认计算图中的操作
        return_elements :frozen_graph_def中的操作名称列表，或其中的张量名称
        """
        return_elements = tf.import_graph_def(frozen_graph_def,
                                              return_elements=return_elements)
    return return_elements

"""
非极大值抑制
去除冗余的bbox
Argument: 
    bboxes:经过post_process_boxes处理过的bbox
    iou_threshold:iou的阈值
    sigma:
    method:非极大值所使用的方法
    
Return: 
    返回经过nms后的bboxes
    将属于同一个类别，且IOU值高于iou_threshold的bbox置0
"""
def nms(bboxes, iou_threshold, sigma=0.3, method='nms'):
    #bbox所属的类别，使用list(set())的原因是set内容无重复
    classes_in_img = list(set(bboxes[:, 5]))

    best_bboxes = []
    #对同属于同一个类别的bbox做如下操作：
    for cls in classes_in_img:
        cls_mask = (bboxes[:, 5] == cls)
        cls_bboxes = bboxes[cls_mask]
        #判断边界框的数目是否大于0
        while len(cls_bboxes) > 0:
            #选出confidence最大的bbox A
            max_ind = np.argmax(cls_bboxes[:, 4])
            #将bbox A 放入best_bbox中
            best_bbox = cls_bboxes[max_ind]
            best_bboxes.append(best_bbox)
            #在cls_bboxes中剔除bbox A
            cls_bboxes = np.concatenate([cls_bboxes[: max_ind], cls_bboxes[max_ind + 1:]])
            #计算bbox A 与剩下的所有bbox的iou值，并剔除那些iou值高于阈值的bbox
            iou = bboxes_iou(best_bbox[np.newaxis, :4], cls_bboxes[:, :4])

            weight = np.ones((len(iou),), dtype=np.float32)

            assert method in ['nms', 'soft-nms']
            #当iou>阈值的时候，将其置0
            if method == 'nms':
                iou_mask = iou > iou_threshold
                #iou大于阈值的其weight=0
                weight[iou_mask] = 0.0

            if method == 'soft-nms':
                weight = np.exp(-(1.0 * iou ** 2 / sigma))
            #剔除cls_bboxes中 bbox与best_bbox的iou大于阈值的bbox
            cls_bboxes[:, 4] = cls_bboxes[:, 4] * weight
            #剔除cls_bboxes中confidence<0的bbox
            score_mask = cls_bboxes[:, 4] > 0.
            cls_bboxes = cls_bboxes[score_mask]

    return best_bboxes

"""
方法用于后处理
后处理步骤
1.将Darknet中的bboxes缩放回在原始图像中的bboxes
2.裁剪超出边界的bboxes
3.去除无效的bboxes(x_min>x_max,或y_min>y_max)
4.去除置信度低的bboxes,即confidence<score_threshold的bboxes

Args:
    pred_bbox:Darknet53输出结果
    org_img_shape:原始输入图像的尺寸，从屏幕中截取图像的尺寸
    input_size:Darknet53输入图像尺寸
    score_threshold:置信度阈值
    
Returns:
    返回经过处理的bbox坐标、置信度、所属类别
    coor:bbox的坐标
    scores:bbox的置信度
    classes:bbox的类别
"""
def postprocess_boxes(pred_bbox, org_img_shape, input_size, score_threshold):
    #有效值范围0~无穷大
    valid_scale=[0, np.inf]
    #转换为np数组类型
    pred_bbox = np.array(pred_bbox)
    #预测bbox的xywh
    pred_xywh = pred_bbox[:, 0:4]
    #预测的最大iou的bbox置信度
    pred_conf = pred_bbox[:, 4]
    #预测的bbox所属类别的概率
    pred_prob = pred_bbox[:, 5:]

    """
    (1) (x, y, w, h) --> (xmin, ymin, xmax, ymax)
    """
    pred_coor = np.concatenate([pred_xywh[:, :2] - pred_xywh[:, 2:] * 0.5,
                                pred_xywh[:, :2] + pred_xywh[:, 2:] * 0.5], axis=-1)

    """
    (2) (xmin, ymin, xmax, ymax) -> (xmin_org, ymin_org, xmax_org, ymax_org)
    将Darknet预测的bbox按比例缩放回原始尺寸的bbox
    """
    #原始图像尺寸
    org_h, org_w = org_img_shape
    #尺寸比例=min(Darknet53输入尺寸/原始图像尺寸)
    resize_ratio = min(input_size / org_w, input_size / org_h)

    dw = (input_size - resize_ratio * org_w) / 2
    dh = (input_size - resize_ratio * org_h) / 2

    pred_coor[:, 0::2] = 1.0 * (pred_coor[:, 0::2] - dw) / resize_ratio
    pred_coor[:, 1::2] = 1.0 * (pred_coor[:, 1::2] - dh) / resize_ratio

    """
    (3) clip some boxes those are out of range
    裁剪超出边界的bboxes
    bboxes的最小位于图像左上顶点，最大位于图像右下顶点向左向右偏移一个像素点
    """
    pred_coor = np.concatenate([np.maximum(pred_coor[:, :2], [0, 0]),
                                np.minimum(pred_coor[:, 2:], [org_w - 1, org_h - 1])], axis=-1)
    #将x_min>x_max,以及y_min>y_max的bbox置0
    invalid_mask = np.logical_or((pred_coor[:, 0] > pred_coor[:, 2]), (pred_coor[:, 1] > pred_coor[:, 3]))
    pred_coor[invalid_mask] = 0

    """
    (4) discard some invalid boxes
    去除无效的bboxes
    """
    bboxes_scale = np.sqrt(np.multiply.reduce(pred_coor[:, 2:4] - pred_coor[:, 0:2], axis=-1))
    scale_mask = np.logical_and((valid_scale[0] < bboxes_scale), (bboxes_scale < valid_scale[1]))

    """
    (5) discard some boxes with low scores
    去除低置信度的bboxes
    """
    #bbox所属的类别
    classes = np.argmax(pred_prob, axis=-1)
    scores = pred_conf * pred_prob[np.arange(len(pred_coor)), classes]
    score_mask = scores > score_threshold
    mask = np.logical_and(scale_mask, score_mask)
    coors, scores, classes = pred_coor[mask], scores[mask], classes[mask]

    return np.concatenate([coors, scores[:, np.newaxis], classes[:, np.newaxis]], axis=-1)



